import os
from sklearn.model_selection import StratifiedKFold

def get_fold_file_lists(brain_root, hipp_root, n_splits):
    """
    读取脑数据文件列表和标签，根据脑数据命名推导海马文件名，
    用 StratifiedKFold 做5折划分，返回每折的训练和验证路径及标签。

    Args:
        brain_root (str): 脑数据根目录（包含 AD, CN, MCI 子文件夹）
        hipp_root (str): 海马数据根目录（对应脑数据结构）

    Returns:
        folds (list): 长度5，元素为dict，包含：
            - train_brain_files, train_hipp_files, train_labels
            - val_brain_files, val_hipp_files, val_labels
    """
    brain_files = []
    labels = []

    class_map = {"AD": 0, "CN": 1, "MCI": 2}

    # 遍历脑数据文件，获取所有路径和标签
    for cls in ["AD", "CN", "MCI"]:
        cls_dir = os.path.join(brain_root, cls)
        for fname in os.listdir(cls_dir):
            if not fname.endswith(".nii.gz"):
                continue
            brain_files.append(os.path.join(cls_dir, fname))
            labels.append(class_map[cls])

    # 推导对应海马文件名的函数
    def get_hipp_file(brain_file):
        base_name = os.path.basename(brain_file)
        hipp_name = base_name.replace(".nii.gz", "_L_Hipp.nii.gz")
        #cls = base_name.split(os.sep)[-2]  # 类别名# 删除这行可能出错的代码
        # 用脑数据文件的类别目录替换为海马目录对应类别
        # 也可以从脑文件路径中提取类别
        cls = brain_file.split(os.sep)[-2]
        return os.path.join(hipp_root, cls, hipp_name)

    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)

    folds = []

    for train_idx, val_idx in skf.split(brain_files, labels):
        train_brain_files = [brain_files[i] for i in train_idx]
        val_brain_files = [brain_files[i] for i in val_idx]

        train_hipp_files = [get_hipp_file(f) for f in train_brain_files]
        val_hipp_files = [get_hipp_file(f) for f in val_brain_files]

        train_labels = [labels[i] for i in train_idx]
        val_labels = [labels[i] for i in val_idx]

        fold = {
            "train_brain_files": train_brain_files,
            "train_hipp_files": train_hipp_files,
            "train_labels": train_labels,
            "val_brain_files": val_brain_files,
            "val_hipp_files": val_hipp_files,
            "val_labels": val_labels
        }
        folds.append(fold)

    return folds
